#include "iter_solver.hpp"

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::IterativeSolver()
    : Solver<idx_t, ksp_t, ksp_t, dof>()
{  }

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::~IterativeSolver()
{  }

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::SetRelTol(double rtol) {
    rel_tol = rtol;
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::SetAbsTol(double atol) {
    abs_tol = atol;
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::SetMaxIter(int max_it) {
    max_iter = max_it;
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::SetPrintLevel(int print_lvl) {
    print_level = print_lvl;
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
int IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::GetNumIterations() const {
    return final_iter;
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
int IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::GetConverged() const {
    return converged;
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
double IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::GetFinalNorm() const {
    return final_norm;
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::SetPreconditioner(Solver<idx_t, ksp_t, pc_calc_t, dof> & pr) {
    // int my_pid; MPI_Comm_rank(MPI_COMM_WORLD, & my_pid);
    // printf("Proc %d SetPrecond\n", my_pid);
    // MPI_Barrier(MPI_COMM_WORLD);
    prec = & pr;
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::SetOperator(const Operator<idx_t, ksp_t, dof> & op) {
    oper = & op;

    this->setup_time = 0.0;
    // 注意，当ksp_t与pc_t的实例化类型不同时，此处会编译报错，需要另外改写：注意是用高精度的完全setup之后再截断
    if (prec) prec->SetOperator(*oper);
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::truncate()
{
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
double IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::Dot(
    const par_Vector<idx_t, ksp_t, dof> & x, const par_Vector<idx_t, ksp_t, dof> & y) const 
{
    double ret;
    if constexpr (sizeof(ksp_t) < 8)// 点积用double及以上的精度
        ret = vec_dot<idx_t, ksp_t, double, dof>(x, y);
    else// 数据本身就是double或者更高
        ret = vec_dot<idx_t, ksp_t, ksp_t , dof>(x, y);
    return ret;
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
double IterativeSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::Norm(
    const par_Vector<idx_t, ksp_t, dof> & x) const {
    ksp_t ret = this->Dot(x, x);
#ifdef __x86_64__
    if constexpr (std::is_same<ksp_t, __float128>::value)
        ret = sqrtq(ret);
    else             
#endif
        ret = sqrt(ret);
    return ret;
}

template class IterativeSolver<int, double, double, double, 1>;
template class IterativeSolver<int, double, double, double, 2>;
template class IterativeSolver<int, double, double, double, 3>;
template class IterativeSolver<int, double, double, double, 4>;

template class IterativeSolver<int, float , double, double, 1>;
template class IterativeSolver<int, float , double, double, 2>;
template class IterativeSolver<int, float , double, double, 3>;
template class IterativeSolver<int, float , double, double, 4>;

template class IterativeSolver<int, float , float , double, 1>;
template class IterativeSolver<int, float , float , double, 2>;
template class IterativeSolver<int, float , float , double, 3>;
template class IterativeSolver<int, float , float , double, 4>;

template class IterativeSolver<int, float , float , float , 1>;
template class IterativeSolver<int, float , float , float , 2>;
template class IterativeSolver<int, float , float , float , 3>;
template class IterativeSolver<int, float , float , float , 4>;
#ifdef USE_FP16
template class IterativeSolver<int, __fp16, float , double, 1>;
template class IterativeSolver<int, __fp16, float , double, 2>;
template class IterativeSolver<int, __fp16, float , double, 3>;
template class IterativeSolver<int, __fp16, float , double, 4>;

template class IterativeSolver<int, __fp16, float , float , 1>;
template class IterativeSolver<int, __fp16, float , float , 2>;
template class IterativeSolver<int, __fp16, float , float , 3>;
template class IterativeSolver<int, __fp16, float , float , 4>;
#endif

template<typename idx_t, typename src_t, typename dst_t, int dof>
void VecTransPrec(const par_Vector<idx_t, src_t, dof> & src_ss_vec, par_Vector<idx_t, dst_t, dof> & dst_ss_vec)
{
    assert(src_ss_vec.glb_nrows == dst_ss_vec.glb_nrows);
    assert(src_ss_vec.beg_row   == dst_ss_vec.beg_row);
    assert(src_ss_vec.end_row   == dst_ss_vec.end_row);

    const src_t * src = src_ss_vec.local_vector->data;
    dst_t       * dst = dst_ss_vec.local_vector->data;
    const long long tot_len = dst_ss_vec.local_vector->tot_len;
    for (long long i = 0; i < tot_len; i++) {
        dst[i] = src[i];
    }
}
template void VecTransPrec(const par_Vector_I32F64_DOF1 &, par_Vector_I32F64_DOF1 &);
template void VecTransPrec(const par_Vector_I32F64_DOF2 &, par_Vector_I32F64_DOF2 &);
template void VecTransPrec(const par_Vector_I32F64_DOF3 &, par_Vector_I32F64_DOF3 &);
template void VecTransPrec(const par_Vector_I32F64_DOF4 &, par_Vector_I32F64_DOF4 &);

template void VecTransPrec(const par_Vector_I32F64_DOF1 &, par_Vector_I32F32_DOF1 &);
template void VecTransPrec(const par_Vector_I32F64_DOF2 &, par_Vector_I32F32_DOF2 &);
template void VecTransPrec(const par_Vector_I32F64_DOF3 &, par_Vector_I32F32_DOF3 &);
template void VecTransPrec(const par_Vector_I32F64_DOF4 &, par_Vector_I32F32_DOF4 &);

template void VecTransPrec(const par_Vector_I32F32_DOF1 &, par_Vector_I32F32_DOF1 &);
template void VecTransPrec(const par_Vector_I32F32_DOF2 &, par_Vector_I32F32_DOF2 &);
template void VecTransPrec(const par_Vector_I32F32_DOF3 &, par_Vector_I32F32_DOF3 &);
template void VecTransPrec(const par_Vector_I32F32_DOF4 &, par_Vector_I32F32_DOF4 &);

template void VecTransPrec(const par_Vector_I32F32_DOF1 &, par_Vector_I32F64_DOF1 &);
template void VecTransPrec(const par_Vector_I32F32_DOF2 &, par_Vector_I32F64_DOF2 &);
template void VecTransPrec(const par_Vector_I32F32_DOF3 &, par_Vector_I32F64_DOF3 &);
template void VecTransPrec(const par_Vector_I32F32_DOF4 &, par_Vector_I32F64_DOF4 &);

template<typename idx_t, typename data_t, typename res_t, int dof>
res_t vec_dot_intra_proc(const par_Vector<idx_t, data_t, dof> & x, const par_Vector<idx_t, data_t, dof> & y)
{
    return seq_vec_dot<idx_t, data_t, res_t, dof>(*(x.local_vector), *(y.local_vector));
}
template double vec_dot_intra_proc(par_Vector<int, double, 1> const &, par_Vector<int, double, 1> const &);
template double vec_dot_intra_proc(par_Vector<int, double, 2> const &, par_Vector<int, double, 2> const &);
template double vec_dot_intra_proc(par_Vector<int, double, 3> const &, par_Vector<int, double, 3> const &);
template double vec_dot_intra_proc(par_Vector<int, double, 4> const &, par_Vector<int, double, 4> const &);
template double vec_dot_intra_proc(par_Vector<int, float , 1> const &, par_Vector<int, float , 1> const &);
template double vec_dot_intra_proc(par_Vector<int, float , 2> const &, par_Vector<int, float , 2> const &);
template double vec_dot_intra_proc(par_Vector<int, float , 3> const &, par_Vector<int, float , 3> const &);
template double vec_dot_intra_proc(par_Vector<int, float , 4> const &, par_Vector<int, float , 4> const &);
